from typing import Tuple
from dataclasses import dataclass
import pandas as pd
import numpy as np

from . import base
from . import feat_utils


@dataclass
class KDJ(base.BarParams, base.BaseTransformer, base.CopyMixin):
    n: int = 9
    new_cols: Tuple[str, str, str] = ('k', 'd', 'j')

    @classmethod
    def trans(cls,
              low: pd.Series,
              high: pd.Series,
              close: pd.Series,
              n: int = 9,):
        # n日最低价
        min_low_n = low.rolling(n).min().fillna(value=low.expanding().min())
        # n日最高
        max_high_n = high.rolling(n).max().fillna(value=low.expanding().max())

        rsv = (close - min_low_n) / (max_high_n - min_low_n) * 100.  # 按收盘价归一化
        k, d = feat_utils.calc_kd(rsv.values)
        k = np.array(k)
        d = np.array(d)
        # 当日J = 3*k - 2*d
        j = 3*k - 2*d
        return k, d, j

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        k, d, j = self.trans(low=data[self.col_low],
                             high=data[self.col_high],
                             close=data[self.col_close],
                             n=self.n)
        kdj = pd.DataFrame(np.c_[k, d, j], columns=self.new_cols, index=data.index)
        return pd.concat([data, kdj], axis=1)

